// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/compute/registry.h>

#include <algorithm>
#include <memory>
#include <mutex>
#include <unordered_map>
#include <utility>

#include <nebula/compute/function.h>
#include <nebula/compute/function_internal.h>
#include <nebula/compute/registry_internal.h>
#include <turbo/utility/status.h>
#include <nebula/version.h>
#include <turbo/log/logging.h>

namespace nebula::compute {

    class FunctionRegistry::FunctionRegistryImpl {
    public:
        explicit FunctionRegistryImpl(FunctionRegistryImpl *parent = nullptr)
                : parent_(parent) {}

        ~FunctionRegistryImpl() {}

        turbo::Status can_add_function(std::shared_ptr<Function> function, bool allow_overwrite) {
            if (parent_ != nullptr) {
                TURBO_RETURN_NOT_OK(parent_->can_add_function(function, allow_overwrite));
            }
            return DoAddFunction(function, allow_overwrite, /*add=*/false);
        }

        turbo::Status add_function(std::shared_ptr<Function> function, bool allow_overwrite) {
            if (parent_ != nullptr) {
                TURBO_RETURN_NOT_OK(parent_->can_add_function(function, allow_overwrite));
            }
            return DoAddFunction(function, allow_overwrite, /*add=*/true);
        }

        turbo::Status can_add_alias(const std::string &target_name, const std::string &source_name) {
            if (parent_ != nullptr) {
                TURBO_RETURN_NOT_OK(parent_->CanAddFunctionName(target_name,
                        /*allow_overwrite=*/false));
            }
            return DoAddAlias(target_name, source_name, /*add=*/false);
        }

        turbo::Status add_alias(const std::string &target_name, const std::string &source_name) {
            if (parent_ != nullptr) {
                TURBO_RETURN_NOT_OK(parent_->CanAddFunctionName(target_name,
                        /*allow_overwrite=*/false));
            }
            return DoAddAlias(target_name, source_name, /*add=*/true);
        }

        turbo::Status can_add_function_options_type(const FunctionOptionsType *options_type,
                                                bool allow_overwrite = false) {
            if (parent_ != nullptr) {
                TURBO_RETURN_NOT_OK(parent_->can_add_function_options_type(options_type, allow_overwrite));
            }
            return DoAddFunctionOptionsType(options_type, allow_overwrite, /*add=*/false);
        }

        turbo::Status add_function_options_type(const FunctionOptionsType *options_type,
                                             bool allow_overwrite = false) {
            if (parent_ != nullptr) {
                TURBO_RETURN_NOT_OK(parent_->can_add_function_options_type(options_type, allow_overwrite));
            }
            return DoAddFunctionOptionsType(options_type, allow_overwrite, /*add=*/true);
        }

        turbo::Result<std::shared_ptr<Function>> get_function(const std::string &name) const {
            auto it = name_to_function_.find(name);
            if (it == name_to_function_.end()) {
                if (parent_ != nullptr) {
                    return parent_->get_function(name);
                }
                return turbo::not_found_error("No function registered with name: ", name);
            }
            return it->second;
        }

        std::vector<std::string> get_function_names() const {
            std::vector<std::string> results;
            if (parent_ != nullptr) {
                results = parent_->get_function_names();
            }
            for (auto it: name_to_function_) {
                results.push_back(it.first);
            }
            std::sort(results.begin(), results.end());
            return results;
        }

        turbo::Result<const FunctionOptionsType *> get_function_options_type(
                const std::string &name) const {
            auto it = name_to_options_type_.find(name);
            if (it == name_to_options_type_.end()) {
                if (parent_ != nullptr) {
                    return parent_->get_function_options_type(name);
                }
                return turbo::not_found_error("No function options type registered with name: ", name);
            }
            return it->second;
        }

        int num_functions() const {
            return (parent_ == nullptr ? 0 : parent_->num_functions()) +
                   static_cast<int>(name_to_function_.size());
        }

        const Function *cast_function() { return cast_function_; }

    private:
        // must not acquire mutex
        turbo::Status CanAddFunctionName(const std::string &name, bool allow_overwrite) {
            if (parent_ != nullptr) {
                TURBO_RETURN_NOT_OK(parent_->CanAddFunctionName(name, allow_overwrite));
            }
            if (!allow_overwrite) {
                auto it = name_to_function_.find(name);
                if (it != name_to_function_.end()) {
                    return turbo::already_exists_error("Already have a function registered with name: ", name);
                }
            }
            return turbo::OkStatus();
        }

        // must not acquire mutex
        turbo::Status CanAddOptionsTypeName(const std::string &name, bool allow_overwrite) {
            if (parent_ != nullptr) {
                TURBO_RETURN_NOT_OK(parent_->CanAddOptionsTypeName(name, allow_overwrite));
            }
            if (!allow_overwrite) {
                auto it = name_to_options_type_.find(name);
                if (it != name_to_options_type_.end()) {
                    return turbo::already_exists_error(
                            "Already have a function options type registered with name: ", name);
                }
            }
            return turbo::OkStatus();
        }

        turbo::Status DoAddFunction(std::shared_ptr<Function> function, bool allow_overwrite,
                                    bool add) {
#ifndef NDEBUG
            // This validates docstrings extensively, so don't waste time on it
            // in release builds.
            TURBO_RETURN_NOT_OK(function->validate());
#endif

            std::lock_guard<std::mutex> mutation_guard(lock_);

            const std::string &name = function->name();
            TURBO_RETURN_NOT_OK(CanAddFunctionName(name, allow_overwrite));
            if (add) {
                name_to_function_[name] = std::move(function);
                if (name == "cast") {
                    cast_function_ = name_to_function_[name].get();
                }
            }
            return turbo::OkStatus();
        }

        turbo::Status DoAddAlias(const std::string &target_name, const std::string &source_name,
                                 bool add) {
            // source name must exist in this registry or the parent
            // check outside mutex, in case get_function leads to mutex acquisition
            TURBO_MOVE_OR_RAISE(auto func, get_function(source_name));

            std::lock_guard<std::mutex> mutation_guard(lock_);

            // target name must be available in this registry and the parent
            TURBO_RETURN_NOT_OK(CanAddFunctionName(target_name, /*allow_overwrite=*/false));
            if (add) {
                name_to_function_[target_name] = func;
            }
            return turbo::OkStatus();
        }

        turbo::Status DoAddFunctionOptionsType(const FunctionOptionsType *options_type,
                                               bool allow_overwrite, bool add) {
            std::lock_guard<std::mutex> mutation_guard(lock_);

            const std::string name = options_type->type_name();
            TURBO_RETURN_NOT_OK(CanAddOptionsTypeName(name, /*allow_overwrite=*/false));
            if (add) {
                name_to_options_type_[options_type->type_name()] = options_type;
            }
            return turbo::OkStatus();
        }

        FunctionRegistryImpl *parent_;
        std::mutex lock_;
        std::unordered_map<std::string, std::shared_ptr<Function>> name_to_function_;
        std::unordered_map<std::string, const FunctionOptionsType *> name_to_options_type_;

        const Function *cast_function_;
    };

    std::unique_ptr<FunctionRegistry> FunctionRegistry::create() {
        return std::unique_ptr<FunctionRegistry>(new FunctionRegistry());
    }

    std::unique_ptr<FunctionRegistry> FunctionRegistry::create(FunctionRegistry *parent) {
        return std::unique_ptr<FunctionRegistry>(new FunctionRegistry(
                new FunctionRegistry::FunctionRegistryImpl(parent->impl_.get())));
    }

    FunctionRegistry::FunctionRegistry() : FunctionRegistry(new FunctionRegistryImpl()) {}

    FunctionRegistry::FunctionRegistry(FunctionRegistryImpl *impl) { impl_.reset(impl); }

    FunctionRegistry::~FunctionRegistry() {}

    turbo::Status FunctionRegistry::can_add_function(std::shared_ptr<Function> function,
                                                   bool allow_overwrite) {
        return impl_->can_add_function(std::move(function), allow_overwrite);
    }

    turbo::Status FunctionRegistry::add_function(std::shared_ptr<Function> function,
                                                bool allow_overwrite) {
        return impl_->add_function(std::move(function), allow_overwrite);
    }

    turbo::Status FunctionRegistry::can_add_alias(const std::string &target_name,
                                                const std::string &source_name) {
        return impl_->can_add_alias(target_name, source_name);
    }

    turbo::Status FunctionRegistry::add_alias(const std::string &target_name,
                                             const std::string &source_name) {
        return impl_->add_alias(target_name, source_name);
    }

    turbo::Status FunctionRegistry::can_add_function_options_type(
            const FunctionOptionsType *options_type, bool allow_overwrite) {
        return impl_->can_add_function_options_type(options_type, allow_overwrite);
    }

    turbo::Status FunctionRegistry::add_function_options_type(const FunctionOptionsType *options_type,
                                                           bool allow_overwrite) {
        return impl_->add_function_options_type(options_type, allow_overwrite);
    }

    turbo::Result<std::shared_ptr<Function>> FunctionRegistry::get_function(
            const std::string &name) const {
        return impl_->get_function(name);
    }

    std::vector<std::string> FunctionRegistry::get_function_names() const {
        return impl_->get_function_names();
    }

    turbo::Result<const FunctionOptionsType *> FunctionRegistry::get_function_options_type(
            const std::string &name) const {
        return impl_->get_function_options_type(name);
    }

    int FunctionRegistry::num_functions() const { return impl_->num_functions(); }

    const Function *FunctionRegistry::cast_function() const { return impl_->cast_function(); }

    namespace internal {

        static std::unique_ptr<FunctionRegistry> create_builtin_registry() {
            auto registry = FunctionRegistry::create();

            // Register core kernels
            register_scalar_cast(registry.get());
            register_dictionary_decode(registry.get());
            register_vector_hash(registry.get());
            register_vector_selection(registry.get());

            register_scalar_options(registry.get());
            register_vector_options(registry.get());
            register_aggregate_options(registry.get());

            // Register additional kernels

            // Scalar functions
            register_scalar_arithmetic(registry.get());
            register_scalar_boolean(registry.get());
            register_scalar_comparison(registry.get());
            register_scalar_if_else(registry.get());
            register_scalar_nested(registry.get());
            register_scalar_random(registry.get());  // Nullary
            register_scalar_round_arithmetic(registry.get());
            register_scalar_set_lookup(registry.get());
            register_scalar_string_ascii(registry.get());
            register_scalar_string_utf8(registry.get());
            register_scalar_temporal_binary(registry.get());
            register_scalar_temporal_unary(registry.get());
            register_scalar_validity(registry.get());

            // Vector functions
            register_vector_array_sort(registry.get());
            register_vector_cumulative_sum(registry.get());
            register_vector_nested(registry.get());
            register_vector_rank(registry.get());
            RegisterVectorReplace(registry.get());
            register_vector_select_k(registry.get());
            register_vector_sort(registry.get());
            register_vector_run_end_encode(registry.get());
            register_vector_run_end_decode(registry.get());
            register_vector_pairwise(registry.get());

            // Aggregate functions
            register_hash_aggregate_basic(registry.get());
            register_scalar_aggregate_basic(registry.get());
            register_scalar_aggregate_mode(registry.get());
            register_scalar_aggregate_quantile(registry.get());
            register_scalar_aggregate_tdigest(registry.get());
            register_scalar_aggregate_variance(registry.get());

            return registry;
        }

    }  // namespace internal

    FunctionRegistry *get_function_registry() {
        static auto g_registry = internal::create_builtin_registry();
        return g_registry.get();
    }

}  // namespace nebula::compute
